import cv2 as cv
import sys
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import numpy as np
import time
import datetime


def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


# img_dir = '../../../../../large_data/CV2/lesson/Day07'
# img_name = 'football_ball.jpg'
# temp_path = os.path.join(img_dir, img_name)

# img_dir = '../../../../../large_data/pic/'
# img_name = 'DSC05039.JPG'
# img_name = 'DSC05022_1.JPG'
# img_name = 'dog.jpg'
# img_name = 'dog_bird.jpg'
# img_name = 'football.jpg'
# img_name = 'messi5.jpg'
img_dir = '../../../../../large_data/pic/watershed/'
img_name = 'water_coins.jpg'
img_path = os.path.join(img_dir, img_name)

spr = 2
spc = 5
spn = 0
plt.figure(figsize=[13, 7])

sep('load')
img = cv.imread(img_path, cv.IMREAD_COLOR)
print('original shape', img.shape)
H, W, CH = img.shape
H2 = H // 2
W2 = W // 2
my_show_img(img, 'original', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('gray')
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
my_show_img(gray, 'gray', cmap='gray')

sep('bin')
ret, bin = cv.threshold(gray, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)
my_show_img(bin, 'bin', cmap='gray')

sep('opening')
kernel = np.ones((3, 3), dtype=np.uint8)
opening = cv.morphologyEx(bin, cv.MORPH_OPEN, kernel, iterations=2)
my_show_img(opening, 'opening', cmap='gray')

sep('sure bg')
sure_bg = cv.dilate(opening, kernel, iterations=3)
my_show_img(sure_bg, 'open => sure_bg', cmap='gray')

sep('sure fg')
dist_trans = cv.distanceTransform(opening, cv.DIST_L2, 5)
my_show_img(dist_trans, 'open => dist trans', cmap='gray')
ret, sure_fg = cv.threshold(dist_trans, 0.7 * dist_trans.max(), 255, cv.THRESH_BINARY)
my_show_img(sure_fg, 'dist trans => sure_fg', cmap='gray')

sep('unknown')
print_numpy_ndarray_info(sure_bg, 'sure_bg')
print_numpy_ndarray_info(sure_fg, 'sure_fg')
sure_fg = sure_fg.astype(np.uint8)
unknown = cv.subtract(sure_bg, sure_fg)
my_show_img(unknown, 'unknown', cmap='gray')

sep('markers')
ret, markers = cv.connectedComponents(sure_fg)
markers += 1
markers[unknown == 255] = 0
my_show_img(markers, 'markers')

sep('watershed')
res = cv.watershed(img, markers)
img[res == -1] = (0, 255, 0)
my_show_img(img, 'after', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
